{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Jupyter environment detected. Enabling Open3D WebVisualizer.\n",
      "[Open3D INFO] WebRTC GUI backend enabled.\n",
      "[Open3D INFO] WebRTCWindowSystem: HTTP handshake server disabled.\n"
     ]
    }
   ],
   "source": [
    "import open3d as o3d \n",
    "import numpy as np \n",
    "from scipy.spatial.transform import Rotation as R"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = o3d.data.PCDPointCloud()\n",
    "pcd = o3d.io.read_point_cloud(dataset.path)\n",
    "assert (pcd.has_normals())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Detect planar patches "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Detected 10 patches\n"
     ]
    }
   ],
   "source": [
    "# using all defaults\n",
    "# This function then returns a list of detected planar patches, represented as geometry::OrientedBoundingBox objects, \n",
    "# with the third column (i.e., z) of R indicating the planar patch normal vector.\n",
    "oboxes = pcd.detect_planar_patches(\n",
    "    normal_variance_threshold_deg=60,\n",
    "    coplanarity_deg=75,\n",
    "    outlier_ratio=0.75,\n",
    "    min_plane_edge_length=0,\n",
    "    min_num_points=0,\n",
    "    search_param=o3d.geometry.KDTreeSearchParamKNN(knn=30))\n",
    "\n",
    "print(\"Detected {} patches\".format(len(oboxes)))\n",
    "\n",
    "normal_vectors = []\n",
    "for obox in oboxes:\n",
    "    normal_vectors.append(obox.R[:, 2])\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_77330/3704614.py:16: UserWarning: Optimal rotation is not uniquely or poorly defined for the given sets of vectors.\n",
      "  rotaion_matrix = R.align_vectors([normal_vectors[i]], [[0, 0, 1]],)[0].as_matrix()\n"
     ]
    }
   ],
   "source": [
    "geometries = []\n",
    "# visualize the patches, bounding boxes, and normals\n",
    "for i, obox in enumerate(oboxes):\n",
    "    mesh = o3d.geometry.TriangleMesh.create_from_oriented_bounding_box(obox, scale=[1, 1, 0.1])\n",
    "    mesh.paint_uniform_color(obox.color)\n",
    "    normal_mesh = o3d.geometry.TriangleMesh.create_arrow(\n",
    "        cylinder_radius=0.02,\n",
    "        cone_radius=0.05,\n",
    "        cylinder_height=0.3,\n",
    "        cone_height=0.1\n",
    "    )\n",
    "    # move mesh to obox center\n",
    "    normal_mesh.translate(obox.get_center())\n",
    "    # rotate mesh to align with obox normal at the obox center\n",
    "    o3d.geometry.get_rotation_matrix_from_axis_angle(normal_vectors[i])\n",
    "    rotaion_matrix = R.align_vectors([normal_vectors[i]], [[0, 0, 1]],)[0].as_matrix()\n",
    "    normal_mesh.rotate(rotaion_matrix, center=obox.get_center())\n",
    "    # set normal arrow color to obox color\n",
    "    normal_mesh.paint_uniform_color(obox.color)\n",
    "    geometries.append(mesh)\n",
    "    geometries.append(obox)\n",
    "    geometries.append(normal_mesh)\n",
    "geometries.append(pcd)\n",
    "\n",
    "# also add coords frame \n",
    "coords = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.5, origin=[0, 0, 0])\n",
    "geometries.append(coords)\n",
    "\n",
    "o3d.visualization.draw_geometries(geometries,\n",
    "                                  zoom=0.62,\n",
    "                                  front=[0.4361, -0.2632, -0.8605],\n",
    "                                  lookat=[2.4947, 1.7728, 1.5541],\n",
    "                                  up=[-0.1726, -0.9630, 0.2071])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plane Segmentation using RANSAC "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plane equation: -0.05x + -0.10y + 0.99z + -1.08 = 0\n"
     ]
    }
   ],
   "source": [
    "plane_model, inliers = pcd.segment_plane(distance_threshold=0.01,\n",
    "                                         ransac_n=3,\n",
    "                                         num_iterations=1000)\n",
    "[a, b, c, d] = plane_model\n",
    "print(f\"Plane equation: {a:.2f}x + {b:.2f}y + {c:.2f}z + {d:.2f} = 0\")\n",
    "\n",
    "inlier_cloud = pcd.select_by_index(inliers)\n",
    "inlier_cloud.paint_uniform_color([1.0, 0, 0])\n",
    "outlier_cloud = pcd.select_by_index(inliers, invert=True)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "o3d.visualization.draw_geometries([inlier_cloud, outlier_cloud],\n",
    "                                  zoom=0.8,\n",
    "                                  front=[-0.4999, -0.1659, -0.8499],\n",
    "                                  lookat=[2.1813, 2.0619, 2.0999],\n",
    "                                  up=[0.1204, -0.9852, 0.1215])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ros_env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
